feat(auth): migrate auth APIs to AmplifyContext - #14836

Nit / suggestion: This ctxOrContextSpec resolution is duplicated verbatim in both server/fetchUserAttributes.ts and server/getCurrentUser.ts. It'd make a nice small shared utility — something like:

// Resolves a server-side argument that may be either the new AmplifyContext
// or a legacy AmplifyServer.ContextSpec, into a concrete AmplifyContext.
export const resolveServerContext = (
	ctxOrContextSpec: AmplifyContext | AmplifyServer.ContextSpec,
): AmplifyContext =>
	'resourcesConfig' in ctxOrContextSpec
		? ctxOrContextSpec
		: getAmplifyServerContext(ctxOrContextSpec).amplify;

Both call sites then become const ctx = resolveServerContext(ctxOrContextSpec);, and the duck-typing discriminator lives in one place if it ever needs to change.

Done in 9b597bb. Extracted the duck-typing into a shared resolveServerContext helper at apis/server/resolveServerContext.ts:

export const resolveServerContext = (
	ctxOrContextSpec: AmplifyContext | AmplifyServer.ContextSpec,
): AmplifyContext => {
	const ctx =
		'resourcesConfig' in ctxOrContextSpec
			? ctxOrContextSpec
			: getAmplifyServerContext(ctxOrContextSpec).amplify;

	return ctx as AmplifyContext;
};

Both server/getCurrentUser.ts and server/fetchUserAttributes.ts now just call const ctx = resolveServerContext(ctxOrContextSpec);, and the discriminator lives in one place. (Kept the cast inside the helper and dropped the now-redundant call-site casts.)

Can you please confirm if this change is intended**

This is the one change that isn't a mechanical context migration, and it looks like it may introduce a behavioral regression:

-		const output =
-			signInInput?.options?.authFlowType === 'USER_AUTH'
-				? await signInWithUserAuth(signInInput)
-				: await signIn(signInInput);
+		const output = await signIn(signInInput);

Two things suggest this drops the cached session during USER_AUTH auto sign-in:

1. Routing through signIn is not equivalent to calling signInWithUserAuth directly — signIn first runs resetAutoSignIn(false) and assertUserNotAuthenticated(ctx). The original code deliberately bypassed those for the USER_AUTH path. Calling them mid-auto-sign-in can disturb the auto-sign-in state machine.
2. In autoSignIn.test.ts, the assertion that the flow forwards session: 'ASDFGHJKL' was silently removed:

 expect(mockHandleUserAuthFlow).toHaveBeenCalledWith(
 	expect.objectContaining({
 		username: user1.username,
-		session: 'ASDFGHJKL',
 	}),
 );

Dropping a coverage assertion to make a test pass is a red flag. Could you confirm whether the cached session is still threaded into the USER_AUTH flow after this change, or whether this is an unintended regression? If the behavior genuinely changed, it likely deserves its own PR + changelog note rather than an inline assertion deletion.

Great catch — you were right, this was an unintended regression, now fixed in 9b597bb.

Confirmed the cause: routing USER_AUTH through signIn() runs resetAutoSignIn(false) first, which dispatches { type: 'RESET' } to autoSignInStore (the false only preserves the callback, not the store state). By the time signInWithUserAuth runs, its if (autoSignInStoreState.active …) guard is false, so the cached session is never forwarded to handleUserAuthFlow — auto-sign-in starts a fresh flow instead of resuming the primed session.

Restored the branch to call signInWithUserAuth directly (matching v7-poc behavior):

const output =
	signInInput?.options?.authFlowType === 'USER_AUTH'
		? await signInWithUserAuth(getGlobalContext(), signInInput)
		: await signIn(signInInput);

I used getGlobalContext() rather than threading ctx as a parameter because auto-sign-in is inherently client-side/post-configure (consistent with the other client-only flows like completeOAuthFlow and handleWebAuthnSignInResult in this split); the v7-poc ctx-threading can land later without a behavior change.

On the dropped assertion — you were right to flag it. I re-added session: 'ASDFGHJKL' in autoSignIn.test.ts, and I also found autoSignInUserConfirmed.test.ts had been changed to assert the regressed behavior (expecting signIn for USER_AUTH); I corrected it to assert signInWithUserAuth(mockCtx, input) is called and signIn is not. Both tests now pass with the fix (full auth suite green: 101/101 suites, 1144 tests).

Question: Do we need to call clearGlobalContext() in afterAll here (and in the other test files that call setGlobalContext)?

It looks like only autoSignIn.test.ts cleans up the global context in afterAll. The other files that call setGlobalContext (this one, signInResumable.test.ts, signInWithCustomAuth.test.ts, etc.) leave it set after the suite finishes, which could leak state across files depending on how much module isolation Jest gives us. Should we standardize on a clearGlobalContext() in afterAll for all of them?

Agreed — standardized in 9b597bb. Added clearGlobalContext() in afterAll to every suite that calls setGlobalContext (this file, signInResumable, signInWithSRP, signInWithCustomAuth, signInWithCustomSRPAuth, signInWithUserAuth, signInWithUserPassword).

While doing this I found the leak was actually being relied on: in signInWithUserPassword.test.ts the Cognito ASF describe had no context setup of its own and depended on the global context leaking from the sibling signIn API happy path cases describe. Adding cleanup exposed that, so I made Cognito ASF self-contained with its own beforeAll(setGlobalContext) / afterAll(clearGlobalContext).

Suggestion: Could createMockAmplifyContext return the mock members already typed as jest.Mock?

Across the migrated tests, this 3-line incantation recurs ~28× (≈14 for mockCtx and ~14 for the per-test customCtx):

(mockCtx.fetchAuthSession as jest.Mock).mockResolvedValue({
	tokens: { accessToken: decodeJWT(mockAccessToken) },
});

The as jest.Mock cast (28×) exists only because the util types these fields as AmplifyContext['fetchAuthSession'] rather than jest.Mock. If createMockAmplifyContext returned a type where fetchAuthSession/getTokens/clearCredentials are jest.Mock (e.g. an intersection like AmplifyContext & { fetchAuthSession: jest.Mock; getTokens: jest.Mock; ... }), the casts disappear. A small convenience helper such as withTokens(ctx, accessToken) could also collapse the repeated mockResolvedValue({ tokens: ... }) setup. Not blocking — just a chance to reduce the boilerplate this migration introduces.

Thanks — this is a good ergonomic improvement and I agree the as jest.Mock casts are noisy. I've deferred it from this PR for now, for a couple of reasons:

1. Scope/risk — this PR is the B-auth split of v6 isolation improvements. #14786 and is meant to be a mechanical context migration plus the regression fix above. Changing the return type of createMockAmplifyContext (e.g. to AmplifyContext & { fetchAuthSession: jest.Mock; getTokens: jest.Mock; clearCredentials: jest.Mock }) touches the shared test util that ~28 call sites depend on, and I'd rather not fold a broad test-ergonomics refactor into the same PR as a behavioral fix.
2. Consistency across splits — the same createMockAmplifyContext util is duplicated in the other B-splits (storage/analytics/notifications each get their own copy per the split plan). I'd prefer to make the typing change once, consistently, rather than diverge auth's copy here.

Happy to do it as a fast follow-up — either a tiny standalone PR against this util, or I can fold it in if you'd prefer it land together. Let me know your preference and I'll take care of it. The withTokens(ctx, accessToken) helper idea is nice too; I'll include that in the same follow-up.

[major] The ctx threading added here is incomplete further down the call chain. getSignInResult (line 576 in this file), which every signIn variant invokes for challenge-to-output mapping, still reads Amplify.getConfig().Auth?.Cognito directly and has no ctx parameter:

export async function getSignInResult(params: { ... }) {
  const authConfig = Amplify.getConfig().Auth?.Cognito; // singleton, not ctx

The concrete gap: a caller that passes signIn(customCtx, input) with a custom userPoolEndpoint will have their InitiateAuth call routed through the custom endpoint (via signInWithSRP(ctx, input) etc.), but if the flow reaches an MFA_SETUP challenge, the associateSoftwareToken call inside getSignInResult uses the singleton endpoint. Both calls happen in the same sign-in flow and may hit different endpoints.

The fix is to add ctx: AmplifyContext to getSignInResult's parameter object and thread it from each signIn variant. If this is deferred to a follow-up split, it is worth tracking explicitly so callers of the fn(ctx, input) overload are not silently getting mixed-context behavior.

Good catch, fixed in 0c91982. You were right — getSignInResult was reading the singleton config, so a custom-endpoint ctx would route InitiateAuth correctly but the MFA_SETUP/associateSoftwareToken path through getSignInResult would hit the singleton endpoint.

Threaded ctx: AmplifyContext as the first param into getSignInResult and switched it to ctx.resourcesConfig.Auth?.Cognito. All 6 callers pass ctx (signInWithSRP/UserAuth/CustomAuth/CustomSRPAuth/UserPassword, confirmSignIn), and the internal WEB_AUTHN recursion (getSignInResult(ctx, result)) plus handleWebAuthnSignInResult were updated too — handleWebAuthnSignInResult now takes ctx as well and uses it for both its authConfig and dispatchSignedInHubEvent(ctx), so the whole challenge-mapping path is single-context now. No deferral needed.

[major] ctx.resourcesConfig is used here for Cognito config, but the credentials and token cleanup later in this function still goes through module-level singletons. Specifically, in the non-OAuth branch:

tokenOrchestrator.clearTokens(); // singleton token store
await clearCredentials();        // clearCredentials imported from @aws-amplify/core, not ctx.clearCredentials()

AmplifyContext exposes ctx.clearCredentials() for exactly this case. Calling the singleton clearCredentials() in a multi-context scenario (which the (ctx, input?) overload is intended for) clears the global Identity pool credential cache, not the per-request one. At minimum, clearCredentials() should be replaced with ctx.clearCredentials(). Whether tokenOrchestrator also needs to be context-scoped depends on this migration phase, but the gap should be explicit if it is a known deferral.

Fixed in 0c91982. Replaced the singleton await clearCredentials() with await ctx.clearCredentials() in the non-OAuth branch and removed the now-unused clearCredentials import from @aws-amplify/core, so per-request credential state is cleared in the (ctx, input?) scenario.

On tokenOrchestrator.clearTokens() — I've left it as the module singleton for now, matching the reference v7-poc which also doesn't context-scope the token orchestrator in this phase. Flagging it explicitly as a known deferral: token-store context-scoping is a broader change beyond this B-auth split.

[minor] The two branches of handleAutoSignInWithCodeOrUserConfirmed pass context inconsistently:

• USER_AUTH path: await signInWithUserAuth(getGlobalContext(), signInInput) -- context passed explicitly
• all other paths (this line): await signIn(signInInput) -- context resolved internally via resolveCtxArgs

Both reach getGlobalContext() at runtime, so there is no functional difference. But the asymmetry is surprising to a reader and looks like an accidental omission. Aligning to await signIn(getGlobalContext(), signInInput) (or a short comment explaining the implicit resolution) would remove the ambiguity.

Fixed in 0c91982. Aligned the else branch to await signIn(getGlobalContext(), signInInput) so both branches pass context explicitly. No functional change (as you noted, both resolve to the global context), just removes the surprising asymmetry.

[minor] The duck-typing discriminator ('resourcesConfig' in ctxOrContextSpec) is the single point of backward compatibility between the new AmplifyContext path and the legacy ContextSpec path, but there are no dedicated unit tests for it.

Both branches are exercised only indirectly through server/getCurrentUser and server/fetchUserAttributes. A small focused test for resolveServerContext itself -- one call with a mock AmplifyContext (has resourcesConfig) and one with a mock ContextSpec (does not) -- would protect the discriminator independently, so a future change to either type fails here rather than silently in a higher-level test.

Added in 0c91982: __tests__/providers/cognito/apis/server/resolveServerContext.test.ts. Two focused cases — one passing a branded AmplifyContext (asserts it's returned as-is and getAmplifyServerContext is not called), and one passing a legacy ContextSpec (asserts it resolves via getAmplifyServerContext(spec).amplify). So the discriminator is now protected independently of the higher-level server API tests.

[nit] This three-line block restates the TypeScript limitation in more words than the cast it documents. A single line captures the same information:

// `'resourcesConfig' in x` doesn't narrow the union; AmplifyClass satisfies AmplifyContext structurally.
return ctx as AmplifyContext;

Done in 0c91982 — collapsed to a single line: // 'resourcesConfig' in x doesn't narrow the union; AmplifyClass satisfies AmplifyContext structurally.